"""
视频文件每帧画面转为蒙版图片存储
"""
#导入相关包
import cv2
import numpy as np

#导入pixellib包
import pixellib
#实例分割
from pixellib.instance import instance_segmentation

#导入弹幕管理模块
from danmu import Danmu_layer,Danmu_track

import os
from PIL import Image
class VideoProcess:
    def __init__(self,videoFile):
        """
        初始化
        """
        self.videoFile = videoFile
    def video2masks(self):
        """
        视频文件转为MASK图片
        """
        #实例分割初始化
        instance = instance_segmentation()
        #加载模型
        instance.load_model('./weights/mask_rcnn_coco.h5')
        
        #筛选类别
        target_classes = instance.select_target_classes(person=True)
        
        #读取视频
        cap = cv2.VideoCapture(self.videoFile)
        #记录帧数
        frame_index = 0
        while True:
            #读取
            ret,frame = cap.read()
            #判断是否读取成功
            if not ret:
                print("视频读取完毕")
                break
            
            #实例分割
            results,output = instance.segmentFrame(frame,segment_target_classes=target_classes)
            #人数
            person_count = len(results['class_ids'])
            if person_count > 0:
                #遮罩
                masks = results['masks']
                #创建黑色底图
                black_bg = np.zeros((frame.shape[:2]))
                
                for p_index in range(person_count):
                    black_bg = np.where(masks[:,:,p_index]==True,255,black_bg)
                #文件名
                mask_file = './masks_img/'+str(frame_index) + '.jpg'
                cv2.imwrite(mask_file,black_bg)#mask_file是要保存文件的路径与文件名，black_bg是要保存的图像数据
                
                print("第%d帧处理完毕"%(frame_index))
            else:
                print("第%d帧没有人"%(frame_index))
                
            frame_index += 1
            #显示渲染图
            cv2.imshow('demo',output)
            #推出条件
            if cv2.waitKey(10) & 0xFF == ord('q'):
                break
        #释放资源
        cap.release()
        cv2.destroyAllWindows()

    def video_composite(self):
        """
        视频合成  弹幕
        1、读取视频第x帧画面
        2、读取第x帧弹幕图面,并用蒙版处理
        3、合成弹幕层与视频层
        4、保存
        """
        #读取视频
        cap = cv2.VideoCapture(self.videoFile)
        
        #获取视频宽度和高度
        frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        #获取视频帧数
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        
        #弹幕层实例化
        danmu_layer = Danmu_layer("danmu_real.txt",frame_w,frame_h)
        
        #记录帧数
        frame_index = 0
        
        #构建视频输出对象
        file_name = './output_video/output.mp4'
        video_writer = cv2.VideoWriter(file_name,cv2.VideoWriter_fourcc(*'mp4v'),fps,(frame_w,frame_h))
        
        while True:
            #读取
            ret,frame = cap.read()
            #判断是否读取成功
            if not ret:
                print("视频读取完毕")
                break
            #获取弹幕层
            frame_danmu_layer = danmu_layer.generate_frame(frame_index)
            
            #对弹幕层进行蒙版处理
            #蒙版文件
            mask_file = './masks_img/'+str(frame_index) + '.jpg'
            
            #判断蒙版文件是否存在
            if os.path.exists(mask_file):
                #读取蒙版文件
                mask_img = cv2.imread(mask_file)
                #转为灰度图
                mask_gray = cv2.cvtColor(mask_img,cv2.COLOR_BGR2GRAY)
                
                #对弹幕层转为numpy数组
                frame_danmu_layer_np = np.array(frame_danmu_layer)
                #对弹幕层alpha通道进行处理
                frame_danmu_layer_np[:,:,3] = np.where(mask_gray==255,0,frame_danmu_layer_np[:,:,3])
                
                #转为RGBA
                frame_danmu_layer = Image.fromarray(frame_danmu_layer_np)
                
            #合成弹幕层与视频层   转为RGBA的Image格式
            frame_rgba = cv2.cvtColor(frame,cv2.COLOR_BGR2RGBA)
            frame_rgba_pil = Image.fromarray(frame_rgba)#使用fromarray()创建成一个对象
            output = Image.alpha_composite(frame_rgba_pil,frame_danmu_layer)#Image.alpha_composite方法传入的参数应该为RGBA图像，不是numpy数组
            
            #转为numpy数组
            output_np = np.asarray(output)
                
            #转为BGR
            output = cv2.cvtColor(output_np,cv2.COLOR_RGBA2BGR)
            
            video_writer.write(output)
                
            #显示渲染图
            cv2.imshow('demo',output)
            frame_index += 1
            #推出条件
            if cv2.waitKey(10) & 0xFF == 27:
                break
        #释放资源
        video_writer.release()
        cap.release()
        cv2.destroyAllWindows()
    
#实例化
vp = VideoProcess('./videos/video.mp4')
vp.video2masks()
vp.video_composite()
